#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EXE文件分包工具 - 简化版
将超过10M的exe文件分包成9M的小包
"""

import os
import hashlib
import json
from pathlib import Path

class EXESplitter:
    def __init__(self):
        self.chunk_size = 9 * 1024 * 1024  # 9MB
        
        # 获取脚本所在目录，然后计算项目根目录
        script_dir = Path(__file__).parent
        project_root = script_dir.parent.parent
        
        self.tools_dir = project_root / "tools"  # tools目录
        self.split_dir = project_root / "split_exes"  # 分包输出目录
        self.restore_dir = project_root / "tools"  # 还原到tools目录
        
    def find_exe_files(self):
        """查找tools目录下的exe文件"""
        exe_files = []
        if self.tools_dir.exists():
            for exe_path in self.tools_dir.glob("*.exe"):
                if exe_path.is_file() and exe_path.stat().st_size > 10 * 1024 * 1024:  # 大于10MB
                    exe_files.append(exe_path)
                    print(f"找到: {exe_path.name} ({exe_path.stat().st_size / 1024 / 1024:.1f}MB)")
        return exe_files
    
    def split_file(self, file_path):
        """分包单个文件"""
        file_size = file_path.stat().st_size
        if file_size <= self.chunk_size:
            print(f"{file_path.name} 小于9MB，跳过")
            return None
            
        # 创建分包目录
        split_subdir = self.split_dir / file_path.stem
        split_subdir.mkdir(parents=True, exist_ok=True)
        
        print(f"分包: {file_path.name}")
        
        # 读取文件并分包
        chunks = []
        with open(file_path, 'rb') as f:
            chunk_num = 1
            while True:
                chunk_data = f.read(self.chunk_size)
                if not chunk_data:
                    break
                    
                chunk_file = split_subdir / f"{file_path.stem}_part{chunk_num:03d}.chunk"
                with open(chunk_file, 'wb') as cf:
                    cf.write(chunk_data)
                
                chunks.append({
                    'file': chunk_file.name,
                    'size': len(chunk_data)
                })
                print(f"  分包 {chunk_num}: {chunk_file.name}")
                chunk_num += 1
        
        # 保存元数据
        metadata = {
            'original': file_path.name,
            'size': file_size,
            'chunks': chunks
        }
        
        with open(split_subdir / "info.json", 'w') as f:
            json.dump(metadata, f, indent=2)
            
        print(f"完成: {len(chunks)} 个分包")
        
        # 删除源文件
        file_path.unlink()
        print(f"已删除源文件: {file_path.name}")
        
        return metadata
    
    def batch_split(self):
        """批量分包"""
        print("开始批量分包...")
        exe_files = self.find_exe_files()
        
        if not exe_files:
            print("没有找到需要分包的exe文件")
            return
            
        results = []
        for exe_file in exe_files:
            result = self.split_file(exe_file)
            if result:
                results.append(result)
        
        # 保存总配置
        config = {
            'files': results,
            'total': len(results)
        }
        
        self.split_dir.mkdir(exist_ok=True)
        with open(self.split_dir / "config.json", 'w') as f:
            json.dump(config, f, indent=2)
            
        print(f"\n分包完成! 共处理 {len(results)} 个文件")
        print(f"分包目录: {self.split_dir}")
    
    def restore_file(self, split_subdir):
        """还原单个文件"""
        info_file = split_subdir / "info.json"
        if not info_file.exists():
            print(f"元数据文件不存在: {info_file}")
            return None
            
        with open(info_file, 'r') as f:
            info = json.load(f)
        
        print(f"还原: {info['original']}")
        
        # 还原文件
        self.restore_dir.mkdir(exist_ok=True)
        output_file = self.restore_dir / info['original']
        
        with open(output_file, 'wb') as out:
            for chunk_info in info['chunks']:
                chunk_file = split_subdir / chunk_info['file']
                with open(chunk_file, 'rb') as cf:
                    out.write(cf.read())
        
        print(f"完成: {output_file}")
        return output_file
    
    def batch_restore(self):
        """批量还原"""
        print("开始批量还原...")
        config_file = self.split_dir / "config.json"
        
        if not config_file.exists():
            print("配置文件不存在")
            return
            
        with open(config_file, 'r') as f:
            config = json.load(f)
        
        restored_files = []
        for file_info in config['files']:
            split_subdir = self.split_dir / Path(file_info['original']).stem
            restored = self.restore_file(split_subdir)
            if restored:
                restored_files.append(restored)
        
        print(f"\n还原完成! 共还原 {len(restored_files)} 个文件")
        print(f"还原目录: {self.restore_dir}")

def main():
    splitter = EXESplitter()
    
    if len(sys.argv) < 2:
        print("用法:")
        print("  python exe_splitter.py split    # 分包")
        print("  python exe_splitter.py restore  # 还原")
        return
    
    action = sys.argv[1]
    
    if action == "split":
        splitter.batch_split()
    elif action == "restore":
        splitter.batch_restore()
    else:
        print("无效操作，使用 split 或 restore")

if __name__ == "__main__":
    import sys
    main() 